import json
def to_tuple(sent):
    triple_list = []
    for triple in sent['triple_list']:
        triple_list.append(tuple(triple))
    sent['triple_list'] = triple_list
def find_head_idx(source, target):
    target_len = len(target)
    for i in range(len(source)):
        if source[i: i + target_len] == target:
            return i
    return -1

if __name__ == '__main__':
    output_path = './NYT/new/train_triples.json'
    train_data_path = './NYT/train_triples.json'
    data = json.load(open(train_data_path))
    for sent in data:
        to_tuple(sent)
    idx = list(range(len(data)))
    for id in idx:
        line = data[id]
        relations = []
        entity = []
        for trip in line['triple_list']:
            entity.append(trip[0])
            entity.append(trip[2])
            relations.append(trip[1])
        F = open(output_path,'a')
        relations_2 = list(set(relations))
        entity_2 = list(set(entity))
        relations_2.sort(key = relations.index)
        entity_2.sort(key = entity.index)
        entity_index = []
        sort_dir = {}
        entity_final=[]
        for ent in entity_2:
            ent_idx = find_head_idx(line['text'], ent)
            entity_index.append(ent_idx)
        for i in range(len(entity_2)):
            sort_dir[i]= entity_index[i]
        sort_dir = sorted(sort_dir.items(), key= lambda x:x[1])
        for key in sort_dir:
            entity_final.append(entity_2[int(key[0])])
        result = json.dumps({
            'text': line['text'],
            'triple_list':line['triple_list'],
            'entity' : [
                    ent for ent in entity_final
                ],
                'relation':[
                    rel for rel in relations_2
                ]

            },ensure_ascii=False, indent=4)
        F.write(result +','+ '\n')
        F.close()
        

        




